#ifndef _CU_BUFFER_H_
#define _CU_BUFFER_H_

#include <stdlib.h>
#include <stdio.h>

#include "..\glew\glew.h"

#include <iostream>

#include <cuda_runtime.h>
#include <cuda_gl_interop.h>

#include "cuMath.h"
#include "cuTypes.h"
#include "cuFunction.h"

#include "../rtCommon/Ray.h"
#include "../glCommon/glBuffer.h"

#define BUFFER_H_TO_D         0x00
#define BUFFER_D_TO_H         0x01

#define BUFFER_READ           0x01          // read  buffer
#define BUFFER_R              0x01          // read  buffer
#define BUFFER_WRITE          0x02          // write buffer
#define BUFFER_W              0x02          // write buffer

#define BUFFER_RIGHT          0x01          // right buffer
#define BUFFER_RT             0x01          // right buffer
#define BUFFER_LEFT           0x02          // left  buffer
#define BUFFER_LT             0x02          // left  buffer

#define BUFFER_READ_AND_WRITE (BUFFER_R  | BUFFER_W)
#define BUFFER_R_AND_W        (BUFFER_R  | BUFFER_W)

#define BUFFER_LEFT_AND_RIGHT (BUFFER_LT | BUFFER_RT)
#define BUFFER_LT_AND_RT      (BUFFER_LT | BUFFER_RT)

#define BUFFER_H              0x01          // allocate buffer host memory
#define BUFFER_HOST           0x01          // allocate buffer host memory

#define BUFFER_D              0x02          // allocate buffer cuda memory
#define BUFFER_DEVICE         0x02          // allocate buffer cuda memory

#define BUFFER_B              0x04          // allocate gl buffer object
#define BUFFER_BUFFER         0x04          // allocate gl buffer object

#define BUFFER_H_AND_D        (BUFFER_H | BUFFER_D)
#define BUFFER_H_AND_B        (BUFFER_H | BUFFER_B)

template <class T> 
class TBuffer 
{
	friend class CglWorld;
private:
	void   size(uint32 size) { m_size   = size; }
	void   offsetR(uint32 offset) { m_offset[m_currR] = offset; }
	void   offsetW(uint32 offset) { m_offset[m_currW] = offset; }
	void   offset(uint8 oFlags, uint32 oR = 0, uint32 oW = 0);

	void   dPtrW(T * ptr) { m_dptr [m_currW] = ptr; }
	void   dPtrR(T * ptr) { m_dptr [m_currR] = ptr; }
	void   dPtr(uint8 dFlags, T  *dptrR = NULL, T  *dptrW = NULL);

	void   vboW(uint32 b) { m_vbo  [m_currW] = b; }
	void   vboR(uint32 b) { m_vbo  [m_currR] = b; }
	void   vbo(uint8 bFlags, uint32 vboR = 0, uint32 vboW = 0);

	void   hPtrW(T * ptr) { m_hptr [m_currW] = ptr; }
	void   hPtrR(T * ptr) { m_hptr [m_currR] = ptr; } 
	void   hPtr(uint8 hFlags, T  *hptrR = NULL, T  *hptrW = NULL);

	void   bDouble(bool bdouble);

public:
  TBuffer();
  ~TBuffer();

  // allocate and Release
  void   Create(uint32 sizeInElems, uint8 cFlags, bool doubleBuffer=false);
  void   Release(uint8 dFlags);

  // Swap buffers for double buffering
  inline void   Swap();

  //////////////////////////////////////////////////////////////////////////
	/// Map : maps buffer object to device memory
	/// IN  : mFlags determine which buffer R, W, or RW to map
	//////////////////////////////////////////////////////////////////////////
  void   Map   (uint8 mFlags  = BUFFER_R);

	//////////////////////////////////////////////////////////////////////////
	/// Unmap : unmaps previously mapped buffer objects
	/// IN    : umFlags determine which buffer R, W, or RW to unmap
	//////////////////////////////////////////////////////////////////////////
  void   Unmap (uint8 umFlags = BUFFER_R);

	//////////////////////////////////////////////////////////////////////////
	/// Update : updates buffer object with data from host memory
	/// IN     : umFlags determine which buffer R, W, or RW to update
	//////////////////////////////////////////////////////////////////////////
  void   Update(uint8 uFlags  = BUFFER_R);

	//////////////////////////////////////////////////////////////////////////
	/// Memcpy : performs copy
	/// IN : dir - can be either H_TO_D or D_TO_H
	/// IN : start - index to start with, default is 0
	/// IN : count - number of elements to copy, default is 0 == size()
	/// IN : cpyFlags determine which buffer R, W, or RW to copy
	//////////////////////////////////////////////////////////////////////////
  void   Memcpy(uint32 dir, 
		            uint32 start    = 0, 
								uint32 count    = 0, 
								uint8  cpyFlags = BUFFER_R);

  void   Memset(T      * pValue   = NULL, 
		            uint32   start    = 0, 
								uint32   count    = 0, 
								uint8    setFlags = BUFFER_R, 
								bool     bUpdateD = true);
 
	//////////////////////////////////////////////////////////////////////////
	/// hPtrW() : returns DEVICE pointer for WRITE array
	//////////////////////////////////////////////////////////////////////////
  inline T     *dPtrW() { return m_dptr [m_currW]; }

	//////////////////////////////////////////////////////////////////////////
	/// hPtrR() : returns DEVICE pointer for READ array
	//////////////////////////////////////////////////////////////////////////
  inline T     *dPtrR() { return m_dptr [m_currR]; }

	//////////////////////////////////////////////////////////////////////////
	/// hPtrLT() : returns DEVICE pointer for LEFT (eq. WRITE) array
	//////////////////////////////////////////////////////////////////////////
	inline T     *dPtrLT() { return m_dptr [m_currW]; }

	//////////////////////////////////////////////////////////////////////////
	/// hPtrRT() : returns DEVICE pointer for RIGHT (eq. READ) array
	//////////////////////////////////////////////////////////////////////////
	inline T     *dPtrRT() { return m_dptr [m_currR]; }


  inline uint32 vboW() { return m_vbo  [m_currW]; }
	inline uint32 vboR() { return m_vbo  [m_currR]; }


	//////////////////////////////////////////////////////////////////////////
	/// hPtrW() : returns HOST pointer for WRITE array
	//////////////////////////////////////////////////////////////////////////
  inline T     *hPtrW() { return m_hptr [m_currW]; }

	//////////////////////////////////////////////////////////////////////////
	/// hPtrR() : returns HOST pointer for READ array
	//////////////////////////////////////////////////////////////////////////
  inline T     *hPtrR() { return m_hptr [m_currR]; }  

	//////////////////////////////////////////////////////////////////////////
	/// hPtrLT() : returns HOST pointer for LEFT  (eq. WRITE) array
	//////////////////////////////////////////////////////////////////////////
	inline T     *hPtrLT() { return m_hptr [m_currW]; }

	//////////////////////////////////////////////////////////////////////////
	/// hPtrRT() : returns HOST pointer for RIGHT (eq. READ)  array
	//////////////////////////////////////////////////////////////////////////
	inline T     *hPtrRT() { return m_hptr [m_currR]; }  

  inline bool   bDouble() { return m_bDouble; }

  inline uint32 size() { return m_size; }

	inline uint32 offsetR() { return m_offset[m_currR]; }
	inline uint32 offsetW() { return m_offset[m_currW]; }

	T      set (T v, uint32 idx, uint32 ib)
	{
		(*this)[ib][idx+m_offset[ib]] = v;
		return v;
	}

	T      get (uint32 idx, uint32 ib)
	{
		return (*this)[ib][idx+m_offset[ib]];
	}

  T    * operator [] (uint32 hostPtrIdx);
  T    * operator () (uint32 cudaPtrIdx);

private:

	void   Null();

  void   CreateD();
  void   CreateB();
  void   CreateH();
 
  void   ReleaseD();
  void   ReleaseB();
  void   ReleaseH();

  uint32 m_size;
	uint32 m_offset[2];

  T     *m_dptr[2];
  T     *m_hptr[2];
  uint32 m_vbo [2];

  bool   m_bDouble;
  uint32 m_currR, m_currW;

	uint8  m_cFlags;
};

template <class T> 
TBuffer<T>::TBuffer()
{
	Null();
}

template <class T>
void TBuffer<T>::Null()
{
	m_size = 0;

	m_offset[0] = 0;
	m_offset[1] = 0;

  m_currR = 0;
	m_currW = 0;

	m_hptr[0] = NULL;
	m_hptr[1] = NULL;

	m_dptr[0] = NULL;
	m_dptr[1] = NULL;

	m_vbo[0] = 0;
	m_vbo[1] = 0;

	m_bDouble = false;

	m_cFlags = 0;
}


template <class T> 
TBuffer<T>::~TBuffer()
{
  Release(m_cFlags);
}

template <class T> 
void TBuffer<T>::Create(uint32 sizeInElems, uint8 cFlags, bool bdouble)
{
	if (sizeInElems == 0) return;

	m_size = sizeInElems;

  m_bDouble = bdouble;

	m_cFlags = cFlags;

  if (m_bDouble) 
  {
    m_currW = 1;
  }

  if (cFlags & BUFFER_H) CreateH();
	if (cFlags & BUFFER_B) CreateB();
	if (cFlags & BUFFER_D) CreateD();
}

template <class T> 
void TBuffer<T>::Release(uint8 rFlags)
{
	if (rFlags & BUFFER_H) ReleaseH();
	if (rFlags & BUFFER_B) ReleaseB();
	if (rFlags & BUFFER_D) ReleaseD();

	m_cFlags = m_cFlags & (~rFlags);
}

template <class T> 
void TBuffer<T>::CreateH()
{
  m_hptr[0] = (T *) new T [m_size];

  if (m_bDouble) m_hptr[1] = (T *) new T [m_size];

#ifdef _DEBUG
	if (m_hptr[0] == NULL) std::cerr << "Host memory allocation failed in TBuffer (read)" << std::endl;

	if (m_bDouble && m_hptr[1] == NULL) std::cerr << "Host memory allocation failed in TBuffer (write)" << std::endl;
#endif
}

template <class T> 
void TBuffer<T>::ReleaseH()
{
  if (m_hptr[0]) 
  {
    delete [] m_hptr[0];
    m_hptr[0] = NULL;
  }

  if (m_hptr[1]) 
  {
    delete [] m_hptr[1];
    m_hptr[1] = NULL;
  }
}

template <class T> 
void TBuffer<T>::CreateD()
{
  cudaMalloc((void **) &m_dptr[0], m_size*sizeof(T));

  if (m_bDouble) 
    cudaMalloc((void **) &m_dptr[1], m_size*sizeof(T));

#ifdef _DEBUG
	if (m_dptr[0] == NULL) std::cerr << "Device memory allocation failed in TBuffer (read)" << std::endl;
	
	if (m_bDouble && m_dptr[1] == NULL) std::cerr << "Device memory allocation failed in TBuffer (write)" << std::endl;
#endif
}

template <class T> 
void TBuffer<T>::ReleaseD()
{
  if (m_dptr[0]) 
  {
    cudaFree(m_dptr[0]);
    m_dptr[0] = NULL;
  }

  if (m_dptr[1]) 
  {
    cudaFree(m_dptr[1]);
    m_dptr[1] = NULL;
  }
}

template <class T> 
void TBuffer<T>::CreateB()
{
	uint32 sizeInBytes = m_size * sizeof(T);

	GL_BUFFER_DESC bufferDesc;
	GL_SUBRESOURCE_DATA subresourceData;
	bufferDesc.ByteWidth = sizeInBytes;

	bool bResult = glCreateBuffer( &bufferDesc, &subresourceData, &m_vbo[0]);

#ifdef _DEBUG
	if (bResult) std::cerr << "GL buffer allocation failed in TBuffer (read)" << std::endl;
#endif

	cudaGLRegisterBufferObject(m_vbo[0]);

  if (m_bDouble) 
	{
    bResult = glCreateBuffer( &bufferDesc, &subresourceData, &m_vbo[1]);

#ifdef _DEBUG
		if (bResult) std::cerr << "GL buffer allocation failed in TBuffer (write)" << std::endl;
#endif

		cudaGLRegisterBufferObject(m_vbo[1]);
	}
}

template <class T> 
void TBuffer<T>::ReleaseB()
{
  if (m_vbo[0]) 
  {
    cudaGLUnregisterBufferObject(m_vbo[0]);
    glDeleteBuffers(1, &m_vbo[0]);
    m_vbo[0] = 0;
  }

  if (m_vbo[1]) 
  {
    cudaGLUnregisterBufferObject(m_vbo[1]);
    glDeleteBuffers(1, &m_vbo[1]);
    m_vbo[1] = 0;
  }
}

template <class T> 
void TBuffer<T>::Swap()
{
  swap_value(m_currR, m_currW);
}

template <class T> 
void TBuffer<T>::Map(uint8 mFlags)
{
  if (mFlags & BUFFER_R && m_vbo[m_currR]) 
		cudaGLMapBufferObject((void **) &m_dptr[m_currR], m_vbo[m_currR]);

	if (mFlags & BUFFER_W && m_vbo[m_currW]) 
		cudaGLMapBufferObject((void **) &m_dptr[m_currW], m_vbo[m_currW]);
}

template <class T> 
void TBuffer<T>::Unmap(uint8 umFlags)
{
  if (umFlags & BUFFER_R && m_vbo[m_currR]) 
		cudaGLUnmapBufferObject(m_vbo[m_currR]); m_dptr[m_currR] = NULL;

	if (umFlags & BUFFER_W && m_vbo[m_currW]) 
		cudaGLUnmapBufferObject(m_vbo[m_currW]); m_dptr[m_currW] = NULL;
}

template <class T>
void TBuffer<T>::Update(uint8 uFlags)
{
	if (uFlags & BUFFER_R && m_vbo[m_currR])
	{
    glBindBuffer(GL_ARRAY_BUFFER, m_vbo[m_currR]);

		//T * hptr = (T *) glMapBuffer(GL_ARRAY_BUFFER, GL_WRITE_ONLY);
		//uint32 oR = m_offset[m_currR];
		//
		//memcpy((void *) &hptr[oR], 
		//	     (void *) &m_hptr[m_currR][oR], 
		//			 m_size * sizeof(T));
    glBufferData(GL_ARRAY_BUFFER, m_size * sizeof(T), m_hptr[m_currR], GL_DYNAMIC_DRAW);

		//glUnmapBuffer(GL_ARRAY_BUFFER);

		glBindBuffer(GL_ARRAY_BUFFER, 0);
	}

	if (uFlags & BUFFER_W && m_vbo[m_currW])
	{
		glBindBuffer(GL_ARRAY_BUFFER, m_vbo[m_currW]);

		//T * hptr = (T *) glMapBuffer(GL_ARRAY_BUFFER, GL_WRITE_ONLY);
		//uint32 oW = m_offset[m_currW];

		//memcpy((void *) &hptr[oW], 
		//	     (void *) &m_hptr[m_currW][oW], 
		//	      m_size * sizeof(T));		
		glBufferData(GL_ARRAY_BUFFER, m_size * sizeof(T), m_hptr[m_currW], GL_DYNAMIC_DRAW);

		//glUnmapBuffer(GL_ARRAY_BUFFER);

		glBindBuffer(GL_ARRAY_BUFFER, 0);
  }
}

template <class T> 
void TBuffer<T>::Memcpy(uint32 dir, uint32 start, uint32 count, uint8 cpyFlags)
{
  if (count==0) 
  {
    count = (uint32) m_size;
  }

	uint32 oR = m_offset[m_currR];
	uint32 oW = m_offset[m_currW];

  switch(dir) 
  {
    case BUFFER_H_TO_D:
			if (cpyFlags & BUFFER_R && m_dptr[m_currR])
				cudaMemcpy( (void *) (m_dptr[m_currR]+oR + start), 
										(void *) (m_hptr[m_currR]+oR + start), 
										count*sizeof(T), cudaMemcpyHostToDevice);
			if (cpyFlags & BUFFER_W && m_dptr[m_currW])
				cudaMemcpy( (void *) (m_dptr[m_currW]+oW + start), 
										(void *) (m_hptr[m_currW]+oW + start), 
										count*sizeof(T), cudaMemcpyHostToDevice);
    break;

    case BUFFER_D_TO_H:
			if (cpyFlags & BUFFER_R && m_dptr[m_currR])
				cudaMemcpy( (void *) (m_hptr[m_currR]+oR + start), 
										(void *) (m_dptr[m_currR]+oR + start), 
										count*sizeof(T), cudaMemcpyDeviceToHost);
			if (cpyFlags & BUFFER_W && m_dptr[m_currW])
				cudaMemcpy( (void *) (m_hptr[m_currW]+oW + start), 
										(void *) (m_dptr[m_currW]+oW + start), 
										count*sizeof(T), cudaMemcpyDeviceToHost);
      break;
    }
}


template <class T> 
void TBuffer<T>::Memset(T      * pValue, 
												uint32   start, 
												uint32   count, 
												uint8    setFlags,
												bool     bUpdateD)
{
	if (count==0) 
	{
		count = (uint32) m_size;
	}

	T v;

	if (pValue == NULL)
		memset(&v, 0, sizeof(v));
	else
		v = *pValue;

	bool bR = setFlags & BUFFER_R;
	bool bW = setFlags & BUFFER_W;

	uint32 oR = m_offset[m_currR];
	uint32 oW = m_offset[m_currW];

  if (m_size <= (start + count))
  {
    for (uint32 idx = 0; idx < count; idx++)
    {
      if (bR) m_hptr[m_currR+oR][idx] = v;
      if (bW) m_hptr[m_currW+oW][idx] = v;
    }
    
		if (bUpdateD)
			Memcpy(BUFFER_H_TO_D, start, count, setFlags);
  }
}

template <class T>
T * TBuffer<T>::operator [](uint32 hostPtrIdx)
{
	T * ptr = NULL;
  switch (hostPtrIdx)
  {
    case BUFFER_R:
      ptr = this->m_hptr[m_currR];
      break;
    case BUFFER_W:
      ptr = this->m_hptr[m_currW];
      break;
    default:
      ptr = NULL;
      break;
  }
	return ptr;
}

template <class T>
T * TBuffer<T>::operator ()(uint32 cudaPtrIdx)
{
	T * ptr = NULL;
  switch (cudaPtrIdx)
  {
    case BUFFER_R:
      ptr = this->m_dptr[m_currR];
      break;
    case BUFFER_W:
      ptr = this->m_dptr[m_currW];
      break;
    default:
      ptr = NULL;
      break;
  }
	return ptr;
}

template <class T>
void TBuffer<T>::offset(uint8 oFlags, uint32 oR, uint32 oW)
{
	if (oFlags & BUFFER_R) m_offset[m_currR] = oR;
	if (oFlags & BUFFER_W) m_offset[m_currW] = oW;
}

template <class T>
void TBuffer<T>::vbo(uint8 bFlags, uint32 bR, uint32 bW)
{
	if (bFlags & BUFFER_R) m_vbo[m_currR] = bR;
	if (bFlags & BUFFER_W) m_vbo[m_currW] = bW;
}

template <class T>
void TBuffer<T>::dPtr(uint8 dFlags, T *dptrR, T *dptrW)
{
	if (dFlags & BUFFER_R) m_dptr[m_currR] = dptrR;
	if (dFlags & BUFFER_W) m_dptr[m_currW] = dptrW;
}

template <class T>
void TBuffer<T>::hPtr(uint8 hFlags, T *hptrR, T *hptrW)
{
	if (hFlags & BUFFER_R) m_hptr[m_currR] = hptrR;
	if (hFlags & BUFFER_W) m_hptr[m_currW] = hptrW;
}

template <class T>
void TBuffer<T>::bDouble(bool bDoubleBuffer)
{
	m_bDouble = bDoubleBuffer;
	m_currR = 0;
	m_currW = bDoubleBuffer ? 1 : 0;
}

 // signed float means regular floating point value
typedef TBuffer<float >           sf_Buffer;
typedef TBuffer<float1>           sf1Buffer;
typedef TBuffer<float2>           sf2Buffer;
typedef TBuffer<float4>           sf4Buffer;

// unsigned floating point means that sign bit is used for some wizardry
typedef TBuffer<float >           uf_Buffer; 
typedef TBuffer<float1>           uf1Buffer;
typedef TBuffer<float2>           uf2Buffer;
typedef TBuffer<float4>           uf4Buffer;

typedef TBuffer<  int >           si_Buffer;
typedef TBuffer<  int1>           si1Buffer;
typedef TBuffer<  int4>           si4Buffer;
typedef TBuffer<  int2>           si2Buffer;


typedef TBuffer< uint >           ui_Buffer;
typedef TBuffer< uint1>           ui1Buffer;
typedef TBuffer< uint2>           ui2Buffer;
typedef TBuffer< uint4>           ui4Buffer;

typedef TBuffer<SRay>           SRayBuffer;

typedef TBuffer<uchar4>          uc4Buffer;

#endif // _CU_BUFFER_H_